GroupBy & Aggregations

Aggregating data is a fundamental task in data analysis. In Spark DataFrames, calling groupBy() partitions the dataset into groups, returning a GroupedData object. You can then apply mathematical aggregates (count, sum, avg, min, max) to resolve the groups into a consolidated DataFrame.

The GroupedData Object

Calling df.groupBy("column_name") does not execute calculations instantly. It prepares Spark to shuffle records so that rows with identical group keys land in the same partition.

Once grouped, you can:

1. Apply a single aggregate directly: df.groupBy("dept").sum("salary")
2. Use agg() to perform multiple, custom aggregates simultaneously (best practice).

Code Example: Basic & Advanced Aggregations

Here is a comprehensive PySpark script executing simple and multi-column aggregations:

from pyspark.sql import SparkSession
from pyspark.sql import functions as F

# 1. Setup Spark
spark = SparkSession.builder \
    .appName("DataFrame Aggregations") \
    .master("local[*]") \
    .getOrCreate()

# 2. Dummy dataset (Sales logs)
sales_data = [
    ("Electronics", "Laptop", 1200, "US"),
    ("Electronics", "Mouse", 40, "US"),
    ("Electronics", "Keyboard", 70, "CA"),
    ("Books", "Fiction Book", 15, "US"),
    ("Books", "Science Manual", 85, "CA"),
    ("Electronics", "Laptop", 1150, "CA")
]
columns = ["category", "product", "revenue", "country"]
df = spark.createDataFrame(sales_data, columns)

# 3. Simple aggregate: Sum revenue by Category
category_revenue_df = df.groupBy("category").sum("revenue")
category_revenue_df.show()

# 4. Advanced Aggregation using agg()
# We will calculate total revenue, average price, and total item count by Category
advanced_summary_df = df.groupBy("category").agg(
    F.sum("revenue").alias("total_revenue"),
    F.round(F.avg("revenue"), 2).alias("avg_revenue"),
    F.count("product").alias("items_sold")
)
advanced_summary_df.show()

# 5. Multi-Key Grouping: Grouping by Category AND Country
multi_group_df = df.groupBy("category", "country").agg(
    F.sum("revenue").alias("revenue")
)
multi_group_df.show()

Key Aggregation Functions Reference

To use these functions, import pyspark.sql.functions: